/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
 * Description: hccl接口类对外头文件
 * Author: gaozhenyue
 * Create: 2024-04-03
 */

#ifndef HCCL_COMM_CPU_PUB_H
#define HCCL_COMM_CPU_PUB_H

#include <vector>
#include <memory>
#include <map>
#include <mutex>
#include "base.h"
#include "hccl_common.h"
#include "mem_device_pub.h"
#include "mem_host_pub.h"
#include "adapter_pub.h"
#include "topoinfo_struct.h"
#include "memory_alloc_ring.h"
#include "hccl_operator.h"
#include "stream_pub.h"
#include "hashtable/universal_concurrent_map.h"
#include "hccl_comm_pub.h"

namespace hccl {
/* * 默认的rank_table, ranklist为空数组;  后续HCCL可以用于判断是否走新分支 */
extern RankTable_t g_hcclDefaultRankTable;

class HcclImplBase;

class HcclCommBase {
public:
    explicit HcclCommBase(u64 inCCLbufferSize = 0, u64 outCCLbufferSize = 0, std::string identifier = "");
    virtual ~HcclCommBase();

    /**********************************************************************
     函 数 名  : hcclComm::init
     功能描述  : 集合通信域初始化
     输入参数  : HcclCommParams& params
             const RankTable_t &rankTable
     输出参数  : 无
     返 回 值  : HcclResult
    **********************************************************************/
    HcclResult init(HcclCommParams &params, const RankTable_t &rankTable = g_hcclDefaultRankTable);
    HcclResult init(HcclCommParams &params, const std::vector<RankInfo> &rankList, WorldGroupInfo &groupCommonData);

    virtual HcclResult InitImpl(DevType deviceType);

    /**********************************************************************
     功能描述  : 创建以group为名字的集合通信
     输入参数  : const std::string& group
             const u32& groupRank
             const std::vector<u32>& groupRanks
     输出参数  : std::shared_ptr<HcclCommBase>& groupComm
     返 回 值  : HcclResult
    **********************************************************************/
    virtual HcclResult CreateGroup(const std::string &group, const u32 &groupRank, const u32 &userRank,
        const std::vector<u32> &groupRanks, std::shared_ptr<HcclCommBase> &groupComm);
    virtual HcclResult CreateGroup(const std::string &group, const u32 &groupRank, const u32 &userRank,
        const std::vector<u32> &groupRanks, std::shared_ptr<hcclComm> &groupComm);

    /**********************************************************************
     功能描述  : 销毁以group为名字的集合通信
     输入参数  : const std::string& group
     返 回 值  : HcclResult
    **********************************************************************/
    HcclResult DestroyGroup(const std::string &group) const;

    /**********************************************************************
     功能描述  : ps侧创建tag所对应的通信链接
     输入参数  : int tag
     返 回 值  : HcclResult
    **********************************************************************/
    HcclResult PsAssociateWorkers(int32_t tag, uint32_t workerRanks[], uint64_t workerNum);

    /**********************************************************************
     功能描述  : 销毁tag所在的通信域
     输入参数  : int tag
     返 回 值  : HcclResult
    **********************************************************************/
    HcclResult DestoryResource(int32_t tag);

    /**********************************************************************
     功能描述  : 生成唯一的集合通信域标识
     输入参数  : 无
     输出参数  : HcclRootInfo* rootInfo
     返 回 值  : HcclResult
    **********************************************************************/
    static HcclResult GetUniqueId(HcclRootInfo *uniqueId);

    HcclResult CsInit(const RoleTableInfo &roleTableInfo);

    HcclResult GetUserRank(u32 &userRank);
    HcclResult GetRankSize(u32 &rankSize);
    void RealeaseBarrierMemory();
    HcclResult GetWorkspaceSubStreamNum(u64 &streamNum, u64 dataSize = 0,
        HcclCMDType optype = HcclCMDType::HCCL_CMD_INVALID) const;
    HcclResult GetWorkspaceMemSize(const std::string &opType, u64 count, HcclDataType dataType,
    u32 &rankSize, u64 &size);
    HcclResult CreateCommCCLbuffer() const;
    HcclResult GetAllReduceScratchSize(const u32 count, const HcclDataType dataType, u64 &scratchSize) const;
    HcclResult SetWorkspaceResource(const std::string &tag, void *memPtr, u64 maxSize,
        std::vector<rtStream_t> &stream);
    HcclResult CreateRemoteOpBasedResources(u64 memSize, const std::string &tag);
    HcclResult CreateOrUpdateRemoteOpBasedResources(u64 memSize, const std::string &tag);
    HcclResult DestroyRemoteOpBasedMem(const std::string &tag);
    std::string GetIdentifier();
    u32 GetRankTableCrc();

    // 目前支持按tag对资源释放、解绑定
    HcclResult ClearOpResource(const std::string &tag);
    // 获取rdma with reduce算子溢出的task信息，然后清除
    HcclResult GetHccsLinkNum(u32 &numHccsLink);
    HcclResult GetDevType(DevType &devType);
    HcclResult IsStandardCard(bool &isStandardCard);
    HcclResult ServiceCancel(ServiceHandle handle);
    HcclResult GetPsMap(std::map<u32, u32> &nodeMap);
    HcclResult SetDevicePid(s32 devicePid);
    HcclResult Break();
    HcclResult GetTagWorkerList(u32 tag, std::vector<u32> &workerList);
    HcclResult GetWorkerList(std::vector<u32> &workerList);
    HcclResult CreateCommResource(const std::string &tag, rtStream_t aiCpuStream, bool isOpbaseMode,
        void **commContext);
    HcclResult GetTopoDesc(HcclTopoDescs *topoDescs, uint32_t topoSize);
    
    void* barrierSendBuf;
    void* barrierRecvBuf;
    std::mutex operatorlock_;
    std::unique_ptr<HcclImplBase> impl_;

    /* * 禁止用户对API类的实体做拷贝构造或拷贝赋值的操作，内部有指针成员变量 */
    HcclCommBase(const HcclCommBase &) = delete;
    HcclCommBase &operator=(const HcclCommBase &) = delete;

    void UpdateIsHaveCpuRank(const RankTable_t &rankTable);
    void UpdateIsHaveCpuRank(const std::vector<RankInfo> &rankList);
    HcclOperator<EmbeddingServiceParam> *EsOpNew();
    void EsOpDelete(HcclOperator<EmbeddingServiceParam> *ptr);
    DeviceMem indirectInCCLbuffer_; /* 保存inCCLbuffer指针的地址 */
    DeviceMem indirectOutCCLbuffer_; /* 保存outCCLbuffer_指针的地址 */
    u64 inCCLbufferSize_;
    u64 outCCLbufferSize_;
    DevType deviceType_;
    DeviceMem barrierInMemory_;
    DeviceMem barrierOutMemory_;
    bool isFirstBarrier_;
    const std::string identifier_;
    bool isHeterogComm_;
    LocklessRingMemoryAllocate<HcclOperator<EmbeddingServiceParam>> lookupRequestPool_;

    // 记录lookup算子接收Request的数量，用以计算dst worker
    UniversalConcurrentMap<s32, u32> lookupRecvRequestCounter_;

    bool isSpecialType_;
    bool isHaveCpuRank_{false};
};
}  // namespace hccl
#endif /* HCCL_COMM_PUB_H */